import argparse

import joblib
import pandas as pd
from more_itertools import chunked

from iflytech_assistant.assistant.dataclasses import FiveSuggestions
from iflytech_assistant.spark_client import instruct_with_profile as instruct

system_prompt = """
我是一个帮你优化日常对话的小助手。

我的主要任务是帮你把话说得更自然、更有吸引力，但还是保持你想表达的意思
"""

user_prompt = """
我在跟我的{target}聊天，我想给对方表达下面这段话，但是不知道怎么说才好。```
{user_input}
```
麻烦帮我按照{tags}，{preferenced_style}的风格改写一下呗。
直接输出给我结果就行了，不需要解释。
"""

parser = argparse.ArgumentParser()
parser.add_argument(
    "-n",
    "--name",
    type=str,
)
args = parser.parse_args()

df = pd.read_excel(".vscode/test.xlsx")

tasks = []
for i, row in df.iterrows():
    user_input = row["input"]
    target = row["target"]
    preferenced_style = row["tags"]
    tags = row["特征"].replace("，", ",").split(",")

    for tag in tags:
        up = user_prompt.format(
            user_input=user_input,
            target=target,
            preferenced_style=preferenced_style,
            tags=tag,
        )
        task = joblib.delayed(instruct)(up, system_prompt)
        tasks.append(task)

results = joblib.Parallel(n_jobs=-1, verbose=10)(tasks)

five_suggestions = []
for responses in chunked(results, 5):
    five_suggestions.append("\n".join([item.llm_response for item in responses]))

df[args.name] = five_suggestions
df.to_excel(".vscode/test.xlsx", index=False)
